/*
 * Copyright (c) 2006-2018, RT-Thread Development Team
 *
 * SPDX-License-Identifier: Apache-2.0
 *
 * Change Logs:
 * Date           Author       Notes
 * 2018-03-22     Murphy       the first version
 */
#include <stdio.h>
#include <stdint.h>
#include <stdlib.h>
#include <rtthread.h>
#include <fal.h>

#include "bl_if.h"
#include "bl_crc32.h"
#include "tiny_aes.h"
#include "quicklz.h"
#include "fastlz.h"

#ifdef TBOOT_USING_HTTP

//#include "webclient.h"

#define DBG_ENABLE
#define DBG_SECTION_NAME          "http_ota"
#ifdef OTA_DOWNLOADER_DEBUG
#define DBG_LEVEL                 DBG_LOG
#else
#define DBG_LEVEL                 DBG_INFO
#endif
#define DBG_COLOR
#include <rtdbg.h>

/* the address offset of download partition */

static int begin_offset = 0;
rt_uint8_t *http_recv_buf = RT_NULL;
rt_size_t http_file_sz = 0, http_recv_sz = 0;

/* handle function, you can store data and so on */
static int http_ota_shard_download_handle(char *buf, int len)
{
    int ret = RT_EOK;

    if ((http_recv_sz + len) > http_file_sz) {
        len = http_file_sz - http_recv_sz;
    }
    rt_memcpy(&http_recv_buf[http_recv_sz], buf, len);

    http_recv_sz += len;

    rt_free(buf);

    return ret;
}

static rt_bool_t http_ota_fw_download(const char *uri)
{
    int ret = RT_EOK, resp_status;
    struct webclient_session* session = RT_NULL;

    UpdateProgress uPState = TB_READ_MORE;
    rt_int32_t numb = 0, rd_cnt = 0, wr_cnt = 0, wr_off = 0;
    bl_header_info_t header = RT_NULL;

    rt_uint8_t *ctx_buf = RT_NULL, *buf = RT_NULL;
    rt_uint32_t crc = 0xFFFFFFFF;
    rt_uint32_t compr_blk_sz = 0;
    rt_uint32_t compr_off = 0;
    rt_size_t decompr_sz = 0;
    rt_size_t blk_sz_off = 0, cpy_len = 0, decrypt_off = 0;
    rt_uint8_t *compr_buf = RT_NULL, *decompr_buf = RT_NULL;
    rt_uint8_t *decrypt_buf = RT_NULL;
#if defined TBOOT_USING_QUICKLZ
    qlz_state_decompress *qlz_decompr_state = RT_NULL;
#endif
#ifdef TBOOT_USING_AES256
	tiny_aes_context *aes_ctx = RT_NULL;
	rt_uint8_t *aes_iv = RT_NULL;
#endif

    /* create webclient session and set header response size */
    session = webclient_session_create(GET_HEADER_BUFSZ);
    if (!session)
    {
        LOG_E("open uri failed.");
        goto _err_exit;
    }
    webclient_set_timeout(session, 1000);

    /* get the real data length */
    webclient_shard_head_function(session, uri, &http_file_sz);

    if (http_file_sz == 0)
    {
        LOG_E("Request file size is 0!");
        goto _err_exit;
    }
    else if (http_file_sz < 0)
    {
        LOG_E("webclient GET request type is chunked.");
        goto _err_exit;
    }

    http_recv_buf = rt_malloc(http_file_sz);
    http_recv_sz = 0;

    /* register the handle function, you can handle data in the function */
    webclient_register_shard_position_function(session, http_ota_shard_download_handle);

    /* the "memory size" that you can provide in the project and uri */
    ret = webclient_shard_position_function(session, uri, begin_offset, http_file_sz, HTTP_OTA_BUFF_LEN);

    /* clear the handle function */
    webclient_register_shard_position_function(session, RT_NULL);

    if (ret == RT_EOK)
    {
        LOG_I("Download file success.");
    }
    else
    {
        goto _err_exit;
    }

    /* verify */
    // read file header
    header = (bl_header_info_t)http_recv_buf;
    numb = sizeof(bl_header_info);
    // check header crc sum
    crc = bl_crc32sum((rt_uint8_t *)header, sizeof(bl_header_info)-4);
    if (header->header_crc != crc) {
        rt_free(header);
        goto _err_exit;
    }
    // check bootloader version
    if ((header->tbl_ver >> 24) > MAJOR_VERSION) {
        rt_free(header);
        goto _err_exit;
    }

    // check magic numb
    if (header->magic_numb != BOOT_MAGIC_NUMB) {
        rt_free(header);
        goto _err_exit;
    }
    // check model
    if (rt_strncmp(header->model, BOOT_INSTRUMENT_MODEL, sizeof(header->model)) != 0) {
        rt_free(header);
        goto _err_exit;
    }
    // check context size
    if (http_recv_sz != header->ctx_size + sizeof(bl_header_info)) {
        rt_free(header);
        goto _err_exit;
    }

    // read all context
    ctx_buf = http_recv_buf + sizeof(bl_header_info);
    numb = http_recv_sz - sizeof(bl_header_info);
    if (numb != header->ctx_size) {
        goto _err_exit;
    }
    // check context crc
    crc = 0xFFFFFFFF;
    crc = bl_crc32(crc, ctx_buf, numb);
    crc ^= 0xFFFFFFFF;
    if (header->ctx_crc != crc) {
        goto _err_exit;
    }

    // malloc read buffer
    buf = rt_malloc(TBOOT_READ_BUFF_SZ);
    if (buf == RT_NULL) {
        goto _err_exit;
    }

    // decryption
    if ((header->flag & BOOT_ENCRYP_MASK) != BOOT_ENCRYP_NONE) {
#ifdef TBOOT_USING_AES256
        if ((header->flag & BOOT_ENCRYP_MASK) == BOOT_ENCRYP_AES256) {
            aes_ctx = rt_malloc(sizeof(tiny_aes_context));
            if (aes_ctx == RT_NULL) {
                goto _err_exit;
            }
            aes_iv = rt_malloc(16);
            if (aes_iv == RT_NULL) {
                goto _err_exit;
            }

            rt_memcpy(aes_iv, BOOT_ENCRYP_AES_KEY, 16);
            tiny_aes_setkey_dec(aes_ctx, (rt_uint8_t *)BOOT_ENCRYP_AES_KEY, 256);
            decrypt_buf = rt_malloc(TBOOT_READ_BUFF_SZ);
        } else
#endif
        {
            goto _err_exit;
        }
    } else {
        decrypt_buf = buf;
    }
    
    if ((header->flag & BOOT_COMPRE_MASK) != BOOT_COMPRE_NONE) {
#if defined TBOOT_USING_FASTLZ
        if ((header->flag & BOOT_COMPRE_MASK) == BOOT_COMPRE_FASTLZ) {
            compr_buf = rt_malloc(TBOOT_COMP_BUFF_SZ);
            decompr_buf = rt_malloc(TBOOT_COMP_BUFF_SZ);
        } else 
#elif defined TBOOT_USING_QUICKLZ
        if ((header->flag & BOOT_COMPRE_MASK) == BOOT_COMPRE_QUICKLZ) {
            qlz_decompr_state = rt_malloc(sizeof(qlz_state_decompress));
            rt_memset(qlz_decompr_state, 0x0, sizeof(qlz_state_decompress));
            compr_buf = rt_malloc(TBOOT_COMP_BUFF_SZ);
            decompr_buf = rt_malloc(TBOOT_COMP_BUFF_SZ);
        } else 
#endif
        {
            goto _err_exit;
        }
    } else {
        decompr_buf = rt_malloc(TBOOT_COMP_BUFF_SZ);
    }

    // erase on chip flash
    dstPartition = fal_partition_find(header->parti_name);
	if (dstPartition == RT_NULL) {
		goto _err_exit;
	}
    if (fal_partition_erase(dstPartition, 0, dstPartition->len) < 0) {
        goto _err_exit;
    }

    rt_thread_yield();
//    crc = 0xFFFFFFFF;
    while (uPState != TB_FIN) {
    switch(uPState){
    case TB_READ_MORE:
#if 0
        numb = dfs_file_read(&fd_src, buf, TBOOT_READ_BUFF_SZ);
        if (numb <= 0) {
            uPState = TB_FIN;
            break;
        } else if (numb != TBOOT_READ_BUFF_SZ) {
//            printf("%d", numb);
        }
#else
        if (rd_cnt == header->ctx_size) {
            uPState = TB_FIN;
            break;
        }
        if (rd_cnt + TBOOT_READ_BUFF_SZ < header->ctx_size) {
            numb = TBOOT_READ_BUFF_SZ;
        } else {
            numb = header->ctx_size - rd_cnt;
        }
        memcpy(buf, &ctx_buf[rd_cnt], numb);
#endif
        rd_cnt += numb;
//        crc = bl_crc32(crc, buf, numb);
    case TB_DECRPT:
        // decryption
#ifdef TBOOT_USING_AES256
        if ((header->flag & BOOT_ENCRYP_MASK) != BOOT_ENCRYP_NONE) {
            tiny_aes_crypt_cbc(aes_ctx, AES_DECRYPT, numb, aes_iv, buf, &decrypt_buf[decrypt_off]);
            decrypt_off += numb;
        } else 
#endif
        {
            decrypt_off += numb;
        }
    case TB_DECOMPR:
        // decompress
        if ((header->flag & BOOT_COMPRE_MASK) != BOOT_COMPRE_NONE) {
            if (decrypt_off < TBOOT_COMP_HEAD_SZ) {
                uPState = TB_READ_MORE;
                break;
            }
            if (compr_off == 0) {
                compr_blk_sz = decrypt_buf[0] << 24 | decrypt_buf[1] << 16 | decrypt_buf[2] << 8 | decrypt_buf[3];
                if (compr_blk_sz > dstPartition->len) {        // error!!!
                    uPState = TB_FIN;
                    break;
                }
                blk_sz_off = TBOOT_COMP_HEAD_SZ;
            } else {
                blk_sz_off = 0;
            }
            if ((compr_blk_sz - compr_off) < (decrypt_off - blk_sz_off)) {
                cpy_len = compr_blk_sz - compr_off;
                rt_memcpy(&compr_buf[compr_off], decrypt_buf+blk_sz_off, cpy_len);
                decrypt_off = decrypt_off - blk_sz_off - cpy_len;
                rt_memmove(decrypt_buf, decrypt_buf+blk_sz_off+cpy_len, decrypt_off);
            } else {
                cpy_len = decrypt_off - blk_sz_off;
                rt_memcpy(&compr_buf[compr_off], decrypt_buf+blk_sz_off, cpy_len);
                decrypt_off = 0;
            }
            compr_off += cpy_len;

            if (compr_off < compr_blk_sz) { // read more block
                uPState = TB_READ_MORE;
                break;
            }

#if defined TBOOT_USING_FASTLZ
            if ((header->flag & BOOT_COMPRE_MASK) == BOOT_COMPRE_FASTLZ) {
                decompr_sz = fastlz_decompress((const void *)compr_buf, compr_off, decompr_buf, TBOOT_READ_BUFF_SZ);
            }
#elif defined TBOOT_USING_QUICKLZ
            if ((header->flag & BOOT_COMPRE_MASK) == BOOT_COMPRE_QUICKLZ) {
                decompr_sz = qlz_decompress((const char *)compr_buf, decompr_buf, qlz_decompr_state);
            }
#endif
            if (decompr_sz <= 0) {        // error!!!
                uPState = TB_FIN;
                break;
            }
        } else {
            if (wr_off + decrypt_off > header->raw_size) {
                decrypt_off = header->raw_size - wr_off;
            }
            rt_memcpy(decompr_buf, decrypt_buf, decrypt_off);
            decompr_sz = decrypt_off;
            decrypt_off = 0;
        }
    case TB_WRITE:
        if (wr_off + decompr_sz > dstPartition->len) {        // error!!!
            uPState = TB_FIN;
            break;
        }
        wr_cnt = fal_partition_write(dstPartition, wr_off, decompr_buf, decompr_sz);
        if (wr_cnt != decompr_sz) {        // error!!!
            uPState = TB_FIN;
            break;
        }
        wr_off += decompr_sz;
        rt_thread_yield();
    case TB_PRE_NEXT:
        if ((header->flag & BOOT_COMPRE_MASK) != BOOT_COMPRE_NONE) {
            compr_off = 0;
            if (decrypt_off > 0) {
                uPState = TB_DECOMPR;
                break;
            }
        }
        uPState = TB_READ_MORE;
    break;
    case TB_FIN:
    break;
    default:
    break;
    }
    }
//    crc ^= 0xFFFFFFFF;

_err_exit:
    if (http_recv_buf != RT_NULL)
        rt_free(http_recv_buf);

    if (session != RT_NULL)
        webclient_close(session);

    if (buf) rt_free(buf);
#ifdef TBOOT_USING_AES256
    if (aes_iv) rt_free(aes_iv);
    if (aes_ctx) rt_free(aes_ctx);
#endif
#if defined TBOOT_USING_QUICKLZ
    if (qlz_decompr_state) rt_free(qlz_decompr_state);
#endif
    if (decrypt_buf) rt_free(decrypt_buf);
    if (compr_buf) rt_free(compr_buf);
    if (decompr_buf) rt_free(decompr_buf);

    if (header->raw_size != wr_off) {
        return RT_FALSE;
    }
    return RT_TRUE;
}

rt_bool_t updateFromHTTP(void)
{
    return http_ota_fw_download(HTTP_OTA_URL);
}
/**
 * msh />http_ota [url]
*/
MSH_CMD_EXPORT(http_ota, Use HTTP to download the firmware);

#endif /* TBOOT_USING_HTTP */
